{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kahneman–Tversky Optimization (KTO) at scale with LoRA\n",
    "\n",
    "This guide provides a step-by-step workflow for preference fine-tuning the [`meta-llama/Meta-Llama-3-8B-Instruct`](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct) model on a multi-GPU Anyscale cluster. You use **LLaMA-Factory** as the training framework and **LoRA** to reduce memory footprint and enable efficient multi-GPU training.\n",
    "\n",
    "**KTO aligns a model to human preferences using **single binary labels (accept or reject)** instead of pairwise “chosen versus rejected” comparisons. KTO directly optimizes the policy on these unary signals, simplifying data preparation while still encouraging preferred behavior and discouraging undesired outputs.\n",
    "\n",
    "## Step 1: Set up your environment\n",
    "\n",
    "### Dependencies\n",
    "First, ensure your environment has the correct libraries. Start with a pre-built container image and install LLaMA-Factory and DeepSpeed on top of it.\n",
    "\n",
    "Recommended container image:\n",
    "```bash\n",
    "anyscale/ray-llm:2.48.0-py311-cu128\n",
    "```\n",
    "\n",
    "Execute the following commands to install the required packages and optional tools for experiment tracking and faster model downloads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Install the specific version of LLaMA-Factory\n",
    "pip install -q llamafactory==0.9.3\n",
    "\n",
    "# (Optional) For accelerated model downloads from Hugging Face\n",
    "pip install -q hf_transfer==0.1.9\n",
    "\n",
    "# (Optional) Acceleration methods (ensure CUDA/Torch compatibility)\n",
    "pip install -q flash-attn==2.8.3 liger-kernel==0.6.2\n",
    "\n",
    "# (Optional) Experiment tracking library\n",
    "pip install -q mlflow==3.4.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model and compute resources\n",
    "\n",
    "| Item | Value |\n",
    "|------|-------|\n",
    "| **Base model** | [`meta-llama/Meta-Llama-3-8B-Instruct`](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct) |\n",
    "| **Workers** | 4 × L40S / A100 (1 GPU each) |\n",
    "\n",
    "Compared to SFT, KTO typically holds two copies of the model (policy and reference), and alignment datasets often use long contexts, so Anyscale recommends GPUs with larger VRAM. Techniques such as **LoRA** and memory-efficient attention can further reduce memory pressure.\n",
    "\n",
    "## Step 2: Prepare the dataset\n",
    "\n",
    "### Understand the dataset\n",
    "This tutorial uses `kto_en_demo`, a unary-preference dataset for KTO. Each record contains a multi-turn ShareGPT-style dialogue with a **binary label** indicating whether the modeled behavior is preferred.\n",
    "\n",
    "This dataset contains:\n",
    "- `messages`: Turn-by-turn chat between a user and the assistant.\n",
    "- `label`: A boolean (`true` or `false`) indicating whether the example is preferred.\n",
    "\n",
    "**Note:** To maintain role alignment in ShareGPT format, you must follow a strict turn order: `human` and `observation` (tool output) must appear in odd-numbered positions, while `gpt` and `function_call` must appear in even-numbered positions. The model learns to generate the content in the `gpt` and `function_call` turns.\n",
    "\n",
    "**Dataset example**\n",
    "```json\n",
    "{\n",
    "\"messages\": [\n",
    "    { \"role\": \"user\", \"content\": \"Compare and contrast the roles of the hippocampus and the prefrontal cortex...\" },\n",
    "    { \"role\": \"assistant\", \"content\": \"The human brain is a highly complex organ, responsible for a myriad of cognitive functions...\" },\n",
    "    { \"role\": \"user\", \"content\": \"Discuss the mechanisms through which the prefrontal cortex ...\" },\n",
    "    { \"role\": \"assistant\", \"content\": \"The prefrontal cortex (PFC)...\" },\n",
    "    { \"role\": \"user\", \"content\": \"Can you elaborate on the role of the amygdala...\" },\n",
    "    { \"role\": \"assistant\", \"content\": \"The amygdala plays a crucial role in the emotional processing of stored memories...\" }\n",
    "],\n",
    "\"label\": true\n",
    "}\n",
    "```\n",
    "\n",
    "### Register the dataset\n",
    "\n",
    "To specify new datasets that are accessible across Ray worker nodes, you must first add a **`dataset_info.json`** to **[storage shared across nodes](https://docs.anyscale.com/configuration/storage#shared)** such as `/mnt/cluster_storage`. This configuration file acts as a central registry for all your datasets. It maps a custom name to your dataset file location, format, and column structure. \n",
    "\n",
    "If you plan to run KTO post-training on the `kto_en_demo` dataset, first complete the setup steps below. Ensure that you place the dataset files in a storage location that all workers can access (for example, a shared mount or object storage). Avoid storing large files on the head node. \n",
    "\n",
    "`dataset_info.json`\n",
    "\n",
    "- `kto_tag` maps the unary preference label used by KTO.\n",
    "- `tags` helps the loader interpret role/content fields in ShareGPT-style records.\n",
    "\n",
    "```json\n",
    "{\n",
    "  \"my_kto_en_demo\": {\n",
    "    \"file_name\": \"/mnt/cluster_storage/kto_en_demo.json\",\n",
    "    \"formatting\": \"sharegpt\",\n",
    "    \"columns\": {\n",
    "      \"messages\": \"messages\",\n",
    "      \"kto_tag\": \"label\"\n",
    "    },\n",
    "    \"tags\": {\n",
    "      \"role_tag\": \"role\",\n",
    "      \"content_tag\": \"content\",\n",
    "      \"user_tag\": \"user\",\n",
    "      \"assistant_tag\": \"assistant\"\n",
    "    }\n",
    "  }\n",
    "}\n",
    "```\n",
    "\n",
    "For a more detailed dataset preparation and formatting guide, see [Choose your data format](https://docs.anyscale.com/llm/fine-tuning/data-preparation#kto)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Make sure all files are accessible to worker nodes\n",
    "# Create a copy of the data in /mnt/cluster_storage\n",
    "wget https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/datasets/sharegpt/kto_en_demo.json -O /mnt/cluster_storage/kto_en_demo.json\n",
    "# Create a copy of the dataset registry in /mnt/cluster_storage\n",
    "cp ../dataset-configs/dataset_info.json /mnt/cluster_storage/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Create the preference-tuning config (KTO and LoRA)\n",
    "\n",
    "Create a YAML file that defines your **KTO** run. It specifies the base model, dataset, **LoRA** settings, KTO hyperparameters, optional acceleration methods, logging, and Ray cluster resources.\n",
    "\n",
    "**Important notes:**\n",
    "- **Acceleration libraries:** You can use `flash_attn` and `liger-kernel` together, but actual speed and memory gains vary with GPU architecture, sequence length, batch size, precision, kernel availability. Benchmark your training workloads to confirm improvements. Note that `fa2` isn't supported on Turing GPUs (e.g., T4).\n",
    "- **Access and paths:** The YAML only needs to be on the **head node**, but any referenced paths (for example, `dataset_dir`, `ray_storage_path`, `output_dir`) must be on **shared storage** (such as `/mnt/cluster_storage/`) visible to all workers.\n",
    "- **Gated models:** If your base model has gated access on Hugging Face, set `HF_TOKEN` in the runtime environment.\n",
    "- **Memory tips:** If VRAM is tight, consider switching to [QLoRA]((https://github.com/ray-project/ray/blob/master/doc/source/ray-overview/examples/llamafactory-llm-fine-tune/notebooks/dpo_qlora.ipynb)) (4/8-bit) and adding the corresponding quantization keys.\n",
    "\n",
    "### Configure LLaMA-Factory with Ray\n",
    "\n",
    "**Note**: To customize the training configuration, edit `train-configs/kto_lora.yaml`. \n",
    "\n",
    "```yaml\n",
    "# kto_lora.yaml\n",
    "\n",
    "### model\n",
    "model_name_or_path: meta-llama/Meta-Llama-3-8B-Instruct\n",
    "trust_remote_code: true\n",
    "\n",
    "### method\n",
    "stage: kto\n",
    "do_train: true\n",
    "finetuning_type: lora\n",
    "lora_rank: 8\n",
    "lora_target: all\n",
    "pref_beta: 0.1\n",
    "\n",
    "### acceleration methods\n",
    "# You can enable both methods at the same time\n",
    "flash_attn: fa2            # Speed up attention and cut activation memory at long context. Use auto on Turing GPUs (e.g., T4)\n",
    "enable_liger_kernel: true  # Reduce VRAM and improve throughput across multiple transformer ops\n",
    "\n",
    "### dataset\n",
    "dataset: my_kto_en_demo\n",
    "dataset_dir: /mnt/cluster_storage\n",
    "\n",
    "template: llama3\n",
    "cutoff_len: 1024\n",
    "max_samples: 1000\n",
    "overwrite_cache: true\n",
    "preprocessing_num_workers: 16\n",
    "\n",
    "### output\n",
    "output_dir: llama3_8b_lora_kto\n",
    "logging_steps: 5\n",
    "save_steps: 50\n",
    "plot_loss: true\n",
    "overwrite_output_dir: true\n",
    "report_to: mlflow   # or none\n",
    "\n",
    "### train\n",
    "per_device_train_batch_size: 1\n",
    "gradient_accumulation_steps: 2\n",
    "num_train_epochs: 3.0  # Low for demo purpose; adjust as needed\n",
    "learning_rate: 5.0e-6\n",
    "bf16: true\n",
    "lr_scheduler_type: cosine\n",
    "warmup_ratio: 0.1\n",
    "ddp_timeout: 180000000\n",
    "\n",
    "### ray\n",
    "ray_run_name: llama3_8b_kto_lora\n",
    "ray_storage_path: /mnt/cluster_storage/\n",
    "ray_num_workers: 4\n",
    "resources_per_worker:\n",
    "  GPU: 1\n",
    "  anyscale/accelerator_shape:4xL40S: 0.001  # Pin a specific node shape\n",
    "  # accelerator_type:L40S: 0.001            # or just request a GPU type\n",
    "\n",
    "ray_init_kwargs:\n",
    "  runtime_env:\n",
    "    env_vars:\n",
    "      # If using gated models like meta-llama/Llama-3-8B-Instruct\n",
    "      HF_TOKEN: <your_huggingface_token>\n",
    "      # Enable faster downloads if hf_transfer is installed:\n",
    "      HF_HUB_ENABLE_HF_TRANSFER: '1'\n",
    "      # If using mlflow for experiments tracking\n",
    "      MLFLOW_TRACKING_URI: \"https://<your_cloud_id>.cloud.databricks.com\"\n",
    "      MLFLOW_TRACKING_TOKEN: \"<mlflow_tracking_token>\"\n",
    "      MLFLOW_EXPERIMENT_NAME: \"/Users/<your_user_id>/experiment_name\"\n",
    "```\n",
    "\n",
    "## Step 4: Train and monitor\n",
    "\n",
    "**Note**: For gated models such as [`meta-llama/Meta-Llama-3-8B-Instruct`](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct), ensure that you accept the license agreement for the models on the Hugging Face site and set `HF_TOKEN` in the runtime environment. If you installed MLflow, configure its credentials. Otherwise, set `report_to: none` in `kto_lora.yaml` to avoid `api_token not set` errors.\n",
    "\n",
    "With all configurations in place, you can launch fine-tuning or post-training in one of two ways:\n",
    "\n",
    "### Option A: Run from a workspace (quick start)\n",
    "\n",
    "The `USE_RAY=1` prefix tells LLaMA-Factory to run in distributed mode on the Ray cluster attached to your workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "USE_RAY=1 llamafactory-cli train ../train-configs/kto_lora.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option B — Run as an Anyscale job (production)\n",
    "\n",
    "For longer or production runs, submit the training as an **Anyscale job**. Jobs run outside your interactive session for better stability, retries, and durable logs. You package LLaMA-Factory and other libraries in a container image and launch with a short job config. See [Run LLaMA-Factory as an Anyscale job](https://docs.anyscale.com/llm/fine-tuning/llamafactory-jobs) for the step-by-step guide.\n",
    "\n",
    "### Tracking with MLflow\n",
    "\n",
    "If you enabled MLflow logging (`report_to: mlflow` in your YAML), LLaMA-Factory logs metrics (loss, learning rate, etc.), parameters, and artifacts to your configured MLflow tracking server.\n",
    "\n",
    "**Example YAML snippet:**\n",
    "\n",
    "```yaml\n",
    "report_to: mlflow\n",
    "\n",
    "ray_init_kwargs:\n",
    "  runtime_env:\n",
    "    env_vars:\n",
    "      MLFLOW_TRACKING_URI: \"https://<your_cloud_id>.cloud.databricks.com\"\n",
    "      MLFLOW_TRACKING_TOKEN: \"<mlflow_tracking_token>\"\n",
    "      MLFLOW_EXPERIMENT_NAME: \"/Users/<your_user_id>/experiment_name\"\n",
    "```\n",
    "\n",
    "**MLFlow example**\n",
    "\n",
    "![MLflow](https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/3.2.3/3.2.3-mlflow.png)\n",
    "\n",
    "For a more detailed guide on tracking experiments with other tools such as Weights and Biases or MLflow, see [Observability and tracking](https://docs.anyscale.com/llm/fine-tuning/observability-and-tracking).\n",
    "\n",
    "## Step 5: Locate checkpoints\n",
    "\n",
    "Ray Train writes checkpoints under `ray_storage_path/ray_run_name`. In this example run, the path is: `/mnt/cluster_storage/llama3_8b_kto_lora`. \n",
    "\n",
    "Inside, you see a **trainer session** directory named like:\n",
    "`TorchTrainer_75e12_00000_0_2025-09-22_17-58-47`.\n",
    "\n",
    "- Ray Train creates `TorchTrainer_*` **when the trainer starts**; the suffix encodes a short run ID and the **start timestamp**.\n",
    "- Within that directory, Ray Train names checkpoints `checkpoint_000xxx/`, where the number is the saved ordered checkpoints.\n",
    "\n",
    "Control the save cadence with `save_strategy` and `save_steps`. For instructions on how to resume interrupted training with `resume_from_checkpoint` and more, see [Understand the artifacts directory](https://docs.anyscale.com/llm/fine-tuning/checkpointing#artifacts-directory).\n",
    "\n",
    "## Step 6: Export the model\n",
    "\n",
    "If you use LoRA, you can keep the base model and adapters separate for [multi-LoRA deployment](https://docs.anyscale.com/llm/serving/multi-lora) or [merge the adapters](https://docs.anyscale.com/llm/fine-tuning/checkpointing#merge-lora) into the base model for low-latency inference. \n",
    "\n",
    "For full fine-tuning or freeze-tuning, export the fine-tuned model directly.\n",
    "\n",
    "You may optionally apply [post-training quantization](https://docs.anyscale.com/llm/fine-tuning/checkpointing#ptq) on merged or full models before serving."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
